import os
import datetime
import shutil
import numpy as np
try:
    from .cvio import cvio
    from .skuCount import gen_label_text
except:
    import sys
    sys.path.append(os.getcwd())
    from function.cvio import cvio
    from skuCount import gen_label_text


def gen_cat_info(labels_path):
    class_name_to_id = {}
    categories = []
    for i, line in enumerate(open(labels_path).readlines()):
        class_id = i# - 1  # starts with -1
        class_name = line.strip().replace(' ','').replace('\n','')
        # if class_id == -1:
        #     assert class_name in ('ignore', '__ignore__')
        #     continue
        if class_name == '':
            continue
        class_name_to_id[class_name] = class_id
        categories.append(dict(
            supercategory=None,
            id=class_id,
            name=class_name,
        ))
    return class_name_to_id, categories


def gen_ann_key_info(points):
    ptarray = np.array(points, dtype=np.float)
    x1 = ptarray[:, 0].min()
    x2 = ptarray[:, 0].max()
    y1 = ptarray[:, 1].min()
    y2 = ptarray[:, 1].max()
    h, w = y2 - y1, x2 - x1
    bbox = [x1, y1, w, h]
    area = (y2 - y1) * (x2 - x1)
    if len(points) <= 2:
        points = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
    return points, bbox, area


def filter_img_ann_list(img_ann_list):
    _img_ann_list = []
    for img, ann in img_ann_list:
        if not (os.path.exists(ann) and os.path.exists(img)):
            continue
        # ann_info = cvio.load_ann(ann)
        # if not len(ann_info['shapes']):
        #     continue
        _img_ann_list.append([img, ann])
    return _img_ann_list


def labelme2coco(input_dir, output_dir, label='', cover_dir=False, copy_img=True, instance=False):

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    else:
        print(output_dir, 'already exists!')
        if cover_dir:
            shutil.rmtree(output_dir)

    data = dict(
        info=dict(
            description=None,
            url=None,
            version=None,
            year=datetime.datetime.now().year,
            contributor=None,
            date_created=datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f'),
        ),
        licenses=[dict(
            url=None,
            id=0,
            name=None,
        )],
        images=[
            # license, url, file_name, height, width, date_captured, id
        ],
        type='instances',
        annotations=[
            # segmentation, area, iscrowd, image_id, bbox, category_id, id
        ],
        categories=[
            # supercategory, id, name
        ],
    )

    if not (label and os.path.exists(label)):
        label = os.path.join(input_dir, 'labels.txt')
        gen_label_text(input_dir, label, instance)

    class_name_to_id, categories = gen_cat_info(label)
    data['categories'].extend(categories)

    ann_out_path = os.path.join(output_dir, 'annotations.json')
    img_out_root = os.path.join(output_dir, 'JPEGImages')

    if copy_img:
        if not os.path.exists(img_out_root):
            os.makedirs(img_out_root)

    img_ann_list = cvio.load_img_ann_list(input_dir, ann_type='json', recursive=False, silent=False)
    img_ann_list = filter_img_ann_list(img_ann_list)
    imgNum = len(img_ann_list)
    if imgNum == 0:
        print('Not found any images at', input_dir)
        return
    tbar = enumerate(img_ann_list)
    ann_id = 0
    for image_id, (imgsrc, annsrc) in tbar:
        img_out_name = os.path.splitext(os.path.basename(imgsrc))[0] + '.jpg'
        img_out_path = os.path.join(img_out_root, img_out_name)

        print('[%d/%d] %s' % (
            image_id + 1, imgNum, img_out_name))

        ann_info = cvio.load_ann(annsrc, ann_type='json')
        # if not len(ann_info['shapes']):
        #     print("Ignore img without any annotations %s" %
        #           os.path.basename(imgsrc))
        #     continue
        height = ann_info['imageHeight']
        width = ann_info['imageWidth']

        if copy_img:
            shutil.copy(imgsrc, img_out_path)
            # img = np.asarray(PIL.Image.open(imgsrc))
            # PIL.Image.fromarray(img).convert('RGB').save(img_out_path)

        data['images'].append(dict(
            license=0,
            url=None,
            file_name='JPEGImages/%s' % (img_out_name),
            height=height,
            width=width,
            date_captured=None,
            id=image_id,
        ))

        for shape in ann_info['shapes']:
            label = shape['label']
            if instance:
                tags = label.split('-')
                if len(tags) and tags[-1].isdigit():
                    tag = '-%s' % tags[-1]
                    label = label[:-len(tag)]
                    shape['label'] = label
            if label not in class_name_to_id:
                print('Ignore label [%s]' % label)
                continue
            points = shape['points']
            cls_id = class_name_to_id[label]
            points, bbox, area = gen_ann_key_info(points)
            if min(bbox[2:]) <= 1:
                print('!!!\n'*30, imgsrc, label, bbox, area,'!!!\n'*30)
                continue
            segm = [np.asarray(points).flatten().tolist()]
            data['annotations'].append(dict(
                id=ann_id,
                image_id=image_id,
                category_id=cls_id,
                segmentation=segm,
                area=area,
                bbox=bbox,
                iscrowd=0,
            ))
            ann_id += 1
    cvio.write_ann(data, ann_out_path)
    print('结果保存至%s.' % output_dir)


if __name__ == '__main__':
    labelme2coco(r'G:\data\datasets\drink\daiding_101\train\zzw_5k_0408',
                           r'G:\data\datasets\drink\daiding_101\drink_coco',
                           r'G:\data\datasets\drink\daiding_101\labels.txt', 
                           cover_dir=False, copy_img=False)
